"""
定时创建验收任务
从 clip_annotated_record 中获取未创建的 clip ，按标注类型和批次组织成验收任务
所有的任务都按照车天进行组织，普通任务最多 10 个 clip 一组，
21 对应的任务最多 100 个 clip 一组，并支持批量通过
"""

from collections import defaultdict
from datetime import datetime

from pymongo import DeleteOne, UpdateMany, InsertOne, UpdateOne
import pytz
from tqdm import tqdm
from connections import mongo_db, mongo_client
from loguru import logger

SH_TZ = pytz.timezone("Asia/Shanghai")

ac_task_col = mongo_db["annotated_acceptance_task"]
check_frame_label_col = mongo_db["check_frame_label"]
clip_anno_col = mongo_db["clip_annotated_record"]


annotate_types = [
    ("parking_free_space", 21),
    ("3D_LANE", 2),
    ("parking_surround_space_detection", 19),
    ("parking_road_sign_mark", 20),
    ("PARKING_3D_LANE", 31),
    ("CITY_4D_LANE", 29),
]


def get_utc_time() -> datetime:
    """
    获取当前UTC时间
    """
    utc_tz = pytz.utc
    now = datetime.now()
    utc_time = now.astimezone(utc_tz)

    return utc_time


def get_priority_by_date(date: str):
    # if date.startswith("2025"):
    #     return 2
    # if date.startswith("2024-11") or date.startswith("2024-12"):
    #     return 3
    return 5


def main():
    # 对每个类别的 clip，按照车天进行分组
    for anno_type in annotate_types:
        logger.info("=" * 20 + f" handle {anno_type}")
        # 1. 找到所有的 clip records
        anno_type_name, anno_type_num = anno_type
        query = {
            "annotate_type": anno_type_num,
            "state": 2,  # 确认
            "accept_status": None,
        }
        if anno_type_num != 21:
            query["is_vis_processed"] = True
        # 只获取关注的部分数据
        projection = {
            "clip_id": 1,
            "name": 1,
            "collect_time": 1,
            "plate_no": 1,
            "annotate_type": 1,
            "frame_cnt": 1,
            "car_type": 1,
            "supplier": 1,
            "confirmed_time": 1,
        }
        cur_type_clip_anno_list = clip_anno_col.find(query, projection).to_list()
        logger.info(f"find total clip_anno: {len(cur_type_clip_anno_list)}")
        if len(cur_type_clip_anno_list) == 0:
            continue

        # 2. 哪些 clip 已经创建了验收任务，将 clip 从中移除
        clip_id_list = [clip["clip_id"] for clip in cur_type_clip_anno_list]
        clip_id_set = set(clip_id_list)
        ac_task_update_ops = []
        exist_ac_tasks = ac_task_col.find(
            {"annotate_type": anno_type_num, "clip_ids": {"$in": clip_id_list}},
            {"clip_ids": 1, "clip_annotated_records": 1},
        ).to_list()
        logger.info(f"exist ac task cnt: {len(exist_ac_tasks)}")
        for ac_task in exist_ac_tasks:
            # 找到交集就是要移除的 clip
            remove_clip_ids = set(ac_task["clip_ids"]) & clip_id_set
            clip_record_map = {
                clip["clip_id"]: clip for clip in ac_task["clip_annotated_records"]
            }

            new_clip_ids, new_clip_records = [], []
            for clip_id in ac_task.clip_ids:
                if clip_id in remove_clip_ids:
                    continue
                new_clip_ids.append(clip_id)
                new_clip_records.append(clip_record_map[clip_id])
            if len(new_clip_ids) == 0:
                ac_task_update_ops.append(DeleteOne({"_id": ac_task["_id"]}))
            else:
                ac_task_update_ops.append(
                    UpdateOne(
                        {"_id": ac_task["_id"]},
                        {
                            "$set": {
                                "clip_ids": new_clip_ids,
                                "clip_annotated_records": new_clip_records,
                                "clip_ids_num": len(new_clip_ids),
                            }
                        },
                    ),
                )

        # 3. 创建新的 ac task
        exist_ac_task_cnt = ac_task_col.count_documents(
            {"annotate_type": anno_type_num}
        )
        logger.info(f"exist ac task cnt: {exist_ac_task_cnt}")

        # 按照车天进行分组
        vehicle_day_group = defaultdict(list)
        for clip in cur_type_clip_anno_list:
            collect_date = str(
                datetime.fromtimestamp(clip["collect_time"] / 1e6, tz=SH_TZ).date()
            )
            confirm_date = str(
                datetime.fromtimestamp(clip["confirmed_time"] / 1e3, tz=SH_TZ).date()
            )
            key = (clip["plate_no"], collect_date, confirm_date)
            vehicle_day_group[key].append(clip)

        # 非 21 的话，10个一组，21 的话，100个一组
        CHUNK = 10 if anno_type_num != 21 else 100
        ac_task_insert_ops = []
        clip_anno_update_ops = []
        task_idx = exist_ac_task_cnt
        for veh_day, clips in vehicle_day_group.items():
            logger.info(f"vehicle day: {veh_day}, clip count: {len(clips)}")
            for ix in range(0, len(clips), CHUNK):
                task_idx += 1
                cur_task_clips = clips[ix : ix + CHUNK]
                new_task = {
                    "task_name": f"[{anno_type_num}][{anno_type_name.upper()}][{veh_day[0]}][{veh_day[1]}]验收任务{task_idx}",
                    "task_description": "",
                    "status": 1,
                    "project": 6,
                    "frequency": "1Hz",
                    "annotate_type": anno_type_num,
                    "is_vis_processed": True,
                    "clip_ids": [clip["clip_id"] for clip in cur_task_clips],
                    "clip_ids_num": len(cur_task_clips),
                    "clip_annotated_records": cur_task_clips,
                    "priority": get_priority_by_date(veh_day[1]),
                    "created_at": get_utc_time(),
                    "plate_no": veh_day[0],
                    "collect_date": veh_day[1],
                    "confirm_date": veh_day[2],
                }
                ac_task_insert_ops.append(InsertOne(new_task))
                clip_anno_update_ops.append(
                    UpdateMany(
                        {
                            "clip_id": {"$in": new_task["clip_ids"]},
                            "annotate_type": anno_type_num,
                        },
                        {
                            "$set": {"accept_status": 10, "accept_finished_at": None}
                        },  # set status pending
                    )
                )
        logger.info("ready to insert task ")
        with mongo_client.start_session() as session, session.start_transaction():
            if ac_task_update_ops:
                ret = ac_task_col.bulk_write(
                    ac_task_update_ops, ordered=False, session=session
                )
            ret = ac_task_col.bulk_write(
                ac_task_insert_ops, ordered=False, session=session
            )
            assert ret.inserted_count == len(ac_task_insert_ops)
            ret = clip_anno_col.bulk_write(
                clip_anno_update_ops, ordered=False, session=session
            )
            assert ret.modified_count == len(cur_type_clip_anno_list)
        logger.info(f"insert task success, total: {len(ac_task_insert_ops)}")


if __name__ == "__main__":
    main()
